#!/usr/bin/python
#Author: Jimmy Saw
#Usage: python get_lineage.py blastfile.tbl 50
#50 is max number of top hits to return
#It is supposed to give you the parent phylum of the hit organism from a BLAST result
#Needs to be in tabular format parsed by Blast2table Perl script

import sys
import re
from Bio import Entrez

Entrez.email = "jimmysaw@gmail.com"

blastfile = sys.argv[1]
bf = open(blastfile, "rU")
bfl = bf.readlines()

maxorgs = sys.argv[2]

hit = re.compile('\[(.*)\]')
lin = re.compile('(\w+) \w+.*')
#lin = re.compile('(\w+).*')
acc = re.compile('\w+|(.*)|')
bad1 = re.compile('unculture.*', re.IGNORECASE)
bad2 = re.compile('clone', re.IGNORECASE)
bad3 = re.compile('unclassified', re.IGNORECASE)
bad4 = re.compile('Bacterium.*')

num = len(bfl)

max = 0

if num > int(maxorgs):
    max = int(maxorgs)
else:
    max = num

print "Top ", maxorgs, " hit phyla:"

for line in bfl[0:max]:
    s = line.split('\t')
    evalue = s[1]
    identities = s[3]
    similarities = s[4]
    locus_tag = s[5]
    accession = s[9]
    desc = s[13].rstrip()
    searchterm = ""
#    if hit.search(desc):
#        p = hit.search(desc)
#        searchterm = p.group(1)
#    else:
#        break
    if bad1.search(desc):
        pass
    elif bad2.search(desc):
        pass
    elif bad3.search(desc):
        pass
    elif bad4.match(desc):
        pass
    else:
        p = lin.search(desc)
        searchterm = p.group(1)

        #pattern = rx1.match(searchterm)
        handle = Entrez.esearch(db="Taxonomy", term=searchterm)
        record = Entrez.read(handle)

        taxid = record["IdList"]

        if len(taxid) != 0:

            h = Entrez.efetch(db="Taxonomy", id=taxid, retmode="xml")

            rec = Entrez.parse(h)

            for r in rec:
                lineage = r.get("Lineage")
                l = lineage.split(';')
                #print l[3].lstrip()
                #print "Search term used =", searchterm #pattern.group(1)
                print locus_tag, "Lineage found =", lineage
                #print locus_tag, "Phylum =", l[2].lstrip()
        else:
            print locus_tag, "Can't query NCBI Taxonomy. Taxid not found!"
            #print "Search term used =", searchterm

bf.close()